package rsa

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"

	_ "crypto/md5"
	_ "crypto/sha1"
	_ "crypto/sha256"
)

type PublicKey struct {
	key *rsa.PublicKey
}

func ParsePublicKey(pubkey []byte) (*PublicKey, error) {
	block, _ := pem.Decode(pubkey)
	if block == nil {
		return nil, errors.New("pub key err")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return &PublicKey{key: pub.(*rsa.PublicKey)}, nil
}

func (pub *PublicKey) Encrypt(data []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, pub.key, data)
}

func (pub *PublicKey) Verify(data []byte, sig []byte) bool {
	return pub.VerifyHash(crypto.SHA256, data, sig)
}

func (pub *PublicKey) VerifyHash(h crypto.Hash, data []byte, sig []byte) bool {
	switch h {
	case crypto.MD5, crypto.SHA1, crypto.SHA256:
		hash := h.New()
		hash.Write(data)
		return rsa.VerifyPKCS1v15(pub.key, h, hash.Sum(nil), sig) == nil
	default:
		return false
	}
}

type PrivateKey struct {
	key *rsa.PrivateKey
}

func ParsePrivateKey(privkey []byte) (*PrivateKey, error) {
	block, _ := pem.Decode([]byte(privkey))
	if block == nil {
		return nil, errors.New("priv key err")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return &PrivateKey{key: priv}, nil
}

func (priv *PrivateKey) Decrypt(data []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, priv.key, data)
}

func (priv *PrivateKey) PublicKey() *PublicKey {
	return &PublicKey{key: &priv.key.PublicKey}
}

func (priv *PrivateKey) Sign(data []byte) ([]byte, error) {
	return priv.SignHash(crypto.SHA256, data)
}

func (priv *PrivateKey) SignHash(h crypto.Hash, data []byte) ([]byte, error) {
	switch h {
	case crypto.MD5, crypto.SHA1, crypto.SHA256:
		hash := h.New()
		hash.Write(data)
		return rsa.SignPKCS1v15(rand.Reader, priv.key, h, hash.Sum(nil))
	default:
		return nil, errors.New("unsupported hash")
	}
}

func GenerateKeyPair(bits int) (privkey, pubkey []byte, err error) {
	priv, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, nil, err
	}
	privkey = pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(priv),
	})

	pubByte, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		return nil, nil, err
	}
	pubkey = pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubByte,
	})
	return
}
